/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file amp_update_scale.cpp
 */

#include "kernel_operator.h"

using namespace AscendC;
constexpr uint32_t BlockBytes = 32;

class AmpUpdateScale
{
public:
    __aicore__ inline AmpUpdateScale()
    {
    }

    __aicore__ inline void Init(GM_ADDR current_scale, GM_ADDR growth_tracker, GM_ADDR found_inf, GM_ADDR updated_scale,
                                GM_ADDR updated_growth_tracker, float growthFactor, float backoffFactor, int32_t growthInterval, TPipe *tmpPipe)
    {
        pipe_ = tmpPipe;

        growthFactor_ = growthFactor;
        backoffFactor_ = backoffFactor;
        growthInterval_ = growthInterval;

        currentScaleGm_.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_CURRENT_SCALE *>(current_scale), 1);
        growthTrackerGm_.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_GROWTH_TRACKER *>(growth_tracker), 1);
        foundInfGm_.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_FOUND_INF *>(found_inf), 1);

        updatedScaleGm_.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_CURRENT_SCALE *>(updated_scale), 1);
        updatedGrowthTrackerGm_.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_GROWTH_TRACKER *>(updated_growth_tracker), 1);

        pipe_->InitBuffer(currentScaleBuf_, BlockBytes);
        pipe_->InitBuffer(growthTrackerBuf_, BlockBytes);
        pipe_->InitBuffer(foundInfBuf_, BlockBytes);

        pipe_->InitBuffer(updatedScaleBuf_, BlockBytes);
        pipe_->InitBuffer(updatedGrowthTrackerBuf_, BlockBytes);
    }

    __aicore__ inline void Process()
    {
        Compute(); // 标量计算，只算一次
    }

private:
    __aicore__ inline void Compute()
    {
        LocalTensor<DTYPE_CURRENT_SCALE> currentScaleLocalT = currentScaleBuf_.Get<DTYPE_CURRENT_SCALE>();
        LocalTensor<DTYPE_GROWTH_TRACKER> growthTrackerLocalT = growthTrackerBuf_.Get<DTYPE_GROWTH_TRACKER>();
        LocalTensor<DTYPE_FOUND_INF> foundInfLocalT = foundInfBuf_.Get<DTYPE_FOUND_INF>();

        LocalTensor<DTYPE_CURRENT_SCALE> updatedScaleLocalT = updatedScaleBuf_.Get<DTYPE_CURRENT_SCALE>();
        LocalTensor<DTYPE_GROWTH_TRACKER> updatedGrowthTrackerLocalT = updatedGrowthTrackerBuf_.Get<DTYPE_GROWTH_TRACKER>();

        DataCopyExtParams copyParamsCurrentScale{1, 1 * sizeof(DTYPE_CURRENT_SCALE), 0, 0, 0};
        DataCopyExtParams copyParamsGrowthTracker{1, 1 * sizeof(DTYPE_GROWTH_TRACKER), 0, 0, 0};
        DataCopyPadExtParams<DTYPE_CURRENT_SCALE> padParamsCurrentScale{false, 0, 0, 0};
        DataCopyPadExtParams<DTYPE_GROWTH_TRACKER> padParamsGrowthTracker{false, 0, 0, 0};

        DataCopy(currentScaleLocalT, currentScaleGm_, BlockBytes / sizeof(DTYPE_CURRENT_SCALE));
        DataCopy(growthTrackerLocalT, growthTrackerGm_, BlockBytes / sizeof(DTYPE_GROWTH_TRACKER));
        DataCopy(foundInfLocalT, foundInfGm_, BlockBytes / sizeof(DTYPE_FOUND_INF));

        currentScale_ = currentScaleLocalT.GetValue(0);
        growthTracker_ = growthTrackerLocalT.GetValue(0);
        foundInf_ = foundInfLocalT.GetValue(0);
        // AscendC::printf("【MY】currentScale_:%f, growthTracker_:%d, foundInf_:%f\n", currentScale_, growthTracker_, foundInf_);

        if (foundInf_ >= 1) {
            currentScale_ *= backoffFactor_;
            growthTracker_ = 0;
        } else {
            successful_ = growthTracker_ + 1;
            if (successful_ == growthInterval_) {
                currentScale_ *= growthFactor_;
                growthTracker_ = 0;
            } else {
                growthTracker_ = successful_;
            }
        }
        updatedScaleLocalT.SetValue(0, currentScale_);
        updatedGrowthTrackerLocalT.SetValue(0, growthTracker_);

        DataCopy(updatedScaleGm_, updatedScaleLocalT, BlockBytes / sizeof(DTYPE_CURRENT_SCALE));
        DataCopy(updatedGrowthTrackerGm_, updatedGrowthTrackerLocalT, BlockBytes / sizeof(DTYPE_GROWTH_TRACKER));
    }

private:
    TPipe *pipe_;

    float growthFactor_;
    float backoffFactor_;
    float currentScale_;
    float foundInf_;
    int32_t growthInterval_;
    int32_t growthTracker_;
    int32_t successful_;

    // 两个updated开头的是输出
    GlobalTensor<DTYPE_CURRENT_SCALE> currentScaleGm_, foundInfGm_, updatedScaleGm_;
    GlobalTensor<DTYPE_GROWTH_TRACKER> growthTrackerGm_, updatedGrowthTrackerGm_;

    TBuf<TPosition::VECCALC> growthTrackerBuf_, foundInfBuf_, currentScaleBuf_;
    TBuf<TPosition::VECCALC> updatedGrowthTrackerBuf_, updatedScaleBuf_;
};

extern "C" __global__ __aicore__ void amp_update_scale(GM_ADDR current_scale, GM_ADDR growth_tracker, GM_ADDR found_inf, GM_ADDR updated_scale, GM_ADDR updated_growth_tracker, GM_ADDR workspace, GM_ADDR tiling)
{
    TPipe pipe;
    GET_TILING_DATA(tiling_data, tiling);
    AmpUpdateScale op;
    op.Init(current_scale, growth_tracker, found_inf, updated_scale, updated_growth_tracker, tiling_data.growthFactor, tiling_data.backoffFactor, tiling_data.growthInterval, &pipe);
    op.Process();
}
